{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "正在生成: 曼德勃罗特集 (z = z^2 + c)...\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LinearSegmentedColormap\n",
    "import time\n",
    "from numba import jit\n",
    "\n",
    "plt.rcParams['font.sans-serif'] = ['SimHei']\n",
    "plt.rcParams['axes.unicode_minus'] = False\n",
    "\n",
    "@jit(nopython=True)\n",
    "def compute_fractal(width, height, xmin, xmax, ymin, ymax, max_iter, formula_code, power=2, c_real=0.0, c_imag=0.0):\n",
    "\n",
    "    # 初始化图像数组\n",
    "    image = np.zeros((height, width))\n",
    "\n",
    "    # 计算每个像素的步长\n",
    "    dx = (xmax - xmin) / width\n",
    "    dy = (ymax - ymin) / height\n",
    "\n",
    "    for x in range(width):\n",
    "        for y in range(height):\n",
    "            # 将像素坐标转换为复平面上的点\n",
    "            zx = xmin + x * dx\n",
    "            zy = ymin + y * dy\n",
    "\n",
    "            # 初始化迭代变量\n",
    "            z_real = zx\n",
    "            z_imag = zy\n",
    "            iter = 0\n",
    "\n",
    "            # 迭代计算\n",
    "            while iter < max_iter:\n",
    "                # 计算模的平方\n",
    "                mod_sq = z_real * z_real + z_imag * z_imag\n",
    "\n",
    "                # 如果超出逃逸半径（这里设为4）\n",
    "                if mod_sq > 4.0:\n",
    "                    break\n",
    "\n",
    "                # 根据公式代码选择不同的迭代公式\n",
    "                if formula_code == 1:  # 曼德勃罗特集\n",
    "                    new_real = z_real * z_real - z_imag * z_imag + zx\n",
    "                    new_imag = 2 * z_real * z_imag + zy\n",
    "                elif formula_code == 2:  # 朱利亚集\n",
    "                    new_real = z_real * z_real - z_imag * z_imag + c_real\n",
    "                    new_imag = 2 * z_real * z_imag + c_imag\n",
    "                elif formula_code == 3:  # 燃烧船分形\n",
    "                    new_real = z_real * z_real - z_imag * z_imag + zx\n",
    "                    new_imag = 2 * abs(z_real) * abs(z_imag) + zy\n",
    "                elif formula_code == 4:  # 三角分形\n",
    "                    new_real = z_real * z_real - z_imag * z_imag - abs(z_real) * z_imag + zx\n",
    "                    new_imag = 2 * z_real * z_imag + abs(z_imag) * z_real + zy\n",
    "                elif formula_code == 5:  # 四次幂分形\n",
    "                    new_real = z_real ** 4 - 6 * z_real * z_real * z_imag * z_imag + z_imag ** 4 + zx\n",
    "                    new_imag = 4 * z_real ** 3 * z_imag - 4 * z_real * z_imag ** 3 + zy\n",
    "                elif formula_code == 6:  # 正弦分形\n",
    "                    new_real = np.sin(z_real) * np.cosh(z_imag) + zx\n",
    "                    new_imag = np.cos(z_real) * np.sinh(z_imag) + zy\n",
    "                elif formula_code == 7:  # 指数分形\n",
    "                    new_real = np.exp(z_real) * np.cos(z_imag) + zx\n",
    "                    new_imag = np.exp(z_real) * np.sin(z_imag) + zy\n",
    "                elif formula_code == 8:  # 自定义公式1: z = z^power + c / z\n",
    "                    if mod_sq > 0:\n",
    "                        new_real = (z_real ** power - z_imag ** power) + (zx * z_real + zy * z_imag) / mod_sq\n",
    "                        new_imag = (power * z_real ** (power - 1) * z_imag) + (zy * z_real - zx * z_imag) / mod_sq\n",
    "                    else:\n",
    "                        break\n",
    "                elif formula_code == 9:  # 自定义公式2: z = sin(z^power) + c\n",
    "                    temp_real = z_real ** power - z_imag ** power\n",
    "                    temp_imag = power * z_real * z_imag ** (power - 1)\n",
    "                    new_real = np.sin(temp_real) * np.cosh(temp_imag) + zx\n",
    "                    new_imag = np.cos(temp_real) * np.sinh(temp_imag) + zy\n",
    "                elif formula_code == 10:  # 自定义公式3: z = z^power + log(z) + c\n",
    "                    if mod_sq > 0:\n",
    "                        log_mod = 0.5 * np.log(mod_sq)\n",
    "                        angle = np.arctan2(z_imag, z_real)\n",
    "                        new_real = (z_real ** power - z_imag ** power) + log_mod + zx\n",
    "                        new_imag = (power * z_real * z_imag ** (power - 1)) + angle + zy\n",
    "                    else:\n",
    "                        break\n",
    "                elif formula_code == 11:  # 蝴蝶分形\n",
    "                    if mod_sq > 0:\n",
    "                        # 计算1/z\n",
    "                        inv_z_real = z_real / mod_sq\n",
    "                        inv_z_imag = -z_imag / mod_sq\n",
    "\n",
    "                        # 计算z^2\n",
    "                        z_sq_real = z_real * z_real - z_imag * z_imag\n",
    "                        z_sq_imag = 2 * z_real * z_imag\n",
    "\n",
    "                        # 计算sin(z^2)\n",
    "                        sin_real = np.sin(z_sq_real) * np.cosh(z_sq_imag)\n",
    "                        sin_imag = np.cos(z_sq_real) * np.sinh(z_sq_imag)\n",
    "\n",
    "                        # 计算exp(1/z)\n",
    "                        exp_real = np.exp(inv_z_real) * np.cos(inv_z_imag)\n",
    "                        exp_imag = np.exp(inv_z_real) * np.sin(inv_z_imag)\n",
    "\n",
    "                        # 计算sin(z^2) * exp(1/z)\n",
    "                        new_real = sin_real * exp_real - sin_imag * exp_imag + zx * power\n",
    "                        new_imag = sin_real * exp_imag + sin_imag * exp_real + zy * power\n",
    "                    else:\n",
    "                        break\n",
    "                elif formula_code == 12:  # 测试不同的分形公式\n",
    "                    # 新增的迭代公式：z = z^3 + z^2 + c\n",
    "                    z_sq_real = z_real * z_real - z_imag * z_imag\n",
    "                    z_sq_imag = 2 * z_real * z_imag\n",
    "                    z_cu_real = z_real * z_sq_real - z_imag * z_sq_imag\n",
    "                    z_cu_imag = z_real * z_sq_imag + z_imag * z_sq_real\n",
    "                    new_real = z_cu_real + z_sq_real + zx\n",
    "                    new_imag = z_cu_imag + z_sq_imag + zy\n",
    "                  else:\n",
    "                      break\n",
    "                # 这里添加你的迭代式\n",
    "\n",
    "\n",
    "            # 记录迭代次数\n",
    "            image[y, x] = iter\n",
    "\n",
    "    return image\n",
    "\n",
    "\n",
    "def plot_fractal(image, max_iter, title, cmap='viridis'):\n",
    "    \"\"\"\n",
    "    绘制分形图像\n",
    "    \"\"\"\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.imshow(image, cmap=cmap, extent=(xmin, xmax, ymin, ymax))\n",
    "\n",
    "    # 创建自定义颜色条\n",
    "    norm = plt.Normalize(vmin=0, vmax=max_iter)\n",
    "    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)\n",
    "    sm.set_array([])\n",
    "    cbar = plt.colorbar(sm, fraction=0.046, pad=0.04)\n",
    "    cbar.set_label('迭代次数', rotation=270, labelpad=20)\n",
    "\n",
    "    plt.title(title, fontsize=16)\n",
    "    plt.xlabel('实部', fontsize=14)\n",
    "    plt.ylabel('虚部', fontsize=14)\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "# 分形参数设置\n",
    "width, height = 800, 800  # 图像分辨率\n",
    "max_iter = 100  # 最大迭代次数\n",
    "\n",
    "# 查看区域设置 (xmin, xmax, ymin, ymax)\n",
    "xmin, xmax = -2.0, 2.0\n",
    "ymin, ymax = -2.0, 2.0\n",
    "\n",
    "# 自定义颜色映射\n",
    "colors = [(0, 0, 0.3), (0, 0, 1), (0, 1, 1), (1, 1, 0), (1, 0, 0)]\n",
    "cmap_name = 'custom_cmap'\n",
    "custom_cmap = LinearSegmentedColormap.from_list(cmap_name, colors)\n",
    "\n",
    "# 测试不同的分形公式\n",
    "formulas = [\n",
    "    (1, \"曼德勃罗特集 (z = z^2 + c)\", 'viridis'),\n",
    "    (2, \"朱利亚集 (z = z^2 + c), c = -0.7 + 0.27i\", 'plasma', -0.7, 0.27),\n",
    "    (3, \"燃烧船分形 (z = |Re(z)|^2 - |Im(z)|^2 + c)\", 'inferno'),\n",
    "    (4, \"三角分形 (z = z^2 - |Re(z)|*Im(z) + c)\", 'magma'),\n",
    "    (5, \"四次幂分形 (z = z^4 + c)\", 'viridis'),\n",
    "    (6, \"正弦分形 (z = sin(z) + c)\", 'cividis'),\n",
    "    (7, \"指数分形 (z = exp(z) + c)\", 'hot'),\n",
    "    (8, \"自定义公式1 (z = z^power + c/z), power=3\", custom_cmap, 0, 0, 3),\n",
    "    (9, \"自定义公式2 (z = sin(z^power) + c), power=2\", 'twilight', 0, 0, 2),\n",
    "    (10, \"自定义公式3 (z = z^power + log(z) + c), power=2\", 'rainbow', 0, 0, 2),\n",
    "    (11, \"蝴蝶分形 (z = sin(z^2) * exp(1/z) + c)\", 'seismic', 0.3, 0.5, 2),\n",
    "    (12, \"自定义公式4 (z = z^3 + z^2 + c)\", 'magma'), \n",
    "]\n",
    "\n",
    "# 生成并显示所有分形\n",
    "for formula in formulas:\n",
    "    formula_code, title, cmap = formula[:3]\n",
    "    c_real = formula[3] if len(formula) > 3 else 0.0\n",
    "    c_imag = formula[4] if len(formula) > 4 else 0.0\n",
    "    power = formula[5] if len(formula) > 5 else 2.0\n",
    "\n",
    "    print(f\"正在生成: {title}...\")\n",
    "    start_time = time.time()\n",
    "\n",
    "    # 计算分形\n",
    "    image = compute_fractal(width, height, xmin, xmax, ymin, ymax,\n",
    "                            max_iter, formula_code, power, c_real, c_imag)\n",
    "\n",
    "    # 显示分形\n",
    "    plot_fractal(image, max_iter, title, cmap)\n",
    "\n",
    "    print(f\"完成! 耗时: {time.time() - start_time:.2f}秒\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MyMLenvAllinOne",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
